import numpy as np
import torch
import matplotlib.pyplot as plt

plt.style.use('fast')

x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)
y = x**2 + 0.2 * torch.rand(x.size())


class SquareModel(torch.nn.Module):
    def __init__(self, in_features, hidden_features, out_features):
        super(SquareModel, self).__init__()
        self.hidden_layer = torch.nn.Linear(in_features, hidden_features)
        self.output_layer = torch.nn.Linear(hidden_features, out_features)

    def forward(self, inputs):
        mid = torch.softmax(self.hidden_layer(inputs), 1, torch.float32)
        outputs = self.output_layer(mid)
        return outputs


model = SquareModel(1, 10, 1)

plt.ion()
fig, ax = plt.subplots()
plt.show()

optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
criterion = torch.nn.MSELoss()

for time in range(1000):
    outputs = model(x)
    loss = criterion(outputs, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if time % 50 == 0:
        ax.set_title('Fitting noise function $x^2$ activated by $softmax()$')
        ax.set_xlabel('$x$')
        ax.set_xticks(np.linspace(-1, 1, 5))
        ax.set_ylabel('$x^2$')
        ax.set_yticks(np.linspace(0, 1.2, 6, endpoint=False))
        ax.scatter(x, y)
        line = ax.plot(x.detach().numpy(),
                       outputs.detach().numpy(), color='gold')
        ax.text(0.5, 0, f'Loss={round(loss.item(), 8)}')
        fig.savefig(f'./SquareModel/{time//50}.png')
        plt.pause(0.1)
        ax.clear()

plt.ioff()
